import asyncio
from typing import List, Dict, Optional
from fastapi import WebSocket,WebSocketDisconnect,WebSocketException

class ConnectionManager:
    def __init__(self):
        self.active_connections:List[WebSocket]=[]
        self.client_connections:Dict[str, WebSocket]={}
        self.main_loop: Optional[asyncio.AbstractEventLoop] = None
    
    def set_main_loop(self, loop: asyncio.AbstractEventLoop):
        """设置主事件循环（由FastAPI应用调用）"""
        self.main_loop = loop
    #创建
    async def connect(self,websocket:WebSocket,user_id:str, message: str="测试通道"):
        await websocket.accept()
        self.active_connections.append(websocket)
        self.client_connections[user_id]=websocket
        print(f"WebSocket连接已建立，用户ID: {user_id}")
        try:
            # 保持连接，等待消息
            while True:
                # 这里可以接收客户端消息，如果需要双向通信
                try:
                    data = await asyncio.wait_for(websocket.receive_text(), timeout=300.0)
                    print(f"收到用户{user_id}的消息: {data}")
                except asyncio.TimeoutError:
                    # 超时检查连接是否还活跃
                    continue
        except WebSocketDisconnect:
            self.disconnect(websocket, user_id)
            print(f"WebSocket连接已断开，用户ID: {user_id}")
    #关闭
    def disconnect(self,websocket:WebSocket,user_id:str):
        if websocket in self.active_connections:
            self.active_connections.remove(websocket)
        self.client_connections.pop(user_id,None)
    #发送（异步方法）
    async def send_personal_message(self,user_id:str,message:str):
        websocket = self.client_connections.get(user_id)
        if websocket:
            try:
                await websocket.send_text(message)
            except (WebSocketDisconnect,RuntimeError) as e:
                print(f"发送消息给用户{user_id}失败:{e}")
                if websocket in self.active_connections:
                    self.active_connections.remove(websocket)
                self.client_connections.pop(user_id,None)
    
    #发送（同步包装方法，用于定时任务等同步上下文中调用）
    def send_message(self, user_id: str, message: str):
        """
        同步方法，用于在定时任务等同步上下文中发送WebSocket消息
        """
        try:
            # 首先检查是否有WebSocket连接
            if user_id not in self.client_connections:
                print(f"用户{user_id}未建立WebSocket连接，跳过发送消息")
                return
            
            websocket = self.client_connections.get(user_id)
            if not websocket:
                print(f"用户{user_id}的WebSocket连接不存在")
                return
            
            # 尝试获取或创建事件循环
            try:
                # 尝试获取当前线程的事件循环
                loop = asyncio.get_event_loop()
                if loop.is_closed():
                    raise RuntimeError("事件循环已关闭")
            except RuntimeError:
                # 如果没有事件循环或已关闭，创建新的
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
            
            # 优先使用主事件循环（FastAPI的主循环）
            if self.main_loop and self.main_loop.is_running():
                # 使用线程安全的方式将任务提交到主事件循环
                future = asyncio.run_coroutine_threadsafe(
                    self.send_personal_message(user_id, message),
                    self.main_loop
                )
                # 不等待结果，避免阻塞
                return
            elif loop.is_running():
                # 如果当前线程的事件循环正在运行，使用 create_task
                asyncio.create_task(self.send_personal_message(user_id, message))
                return
            else:
                # 如果事件循环未运行，使用 run_until_complete
                try:
                    loop.run_until_complete(self.send_personal_message(user_id, message))
                finally:
                    # 如果是新创建的事件循环，关闭它
                    if not loop.is_running() and not loop.is_closed():
                        try:
                            loop.close()
                        except:
                            pass
        except Exception as e:
            print(f"发送消息给用户{user_id}失败（同步包装）: {e}")
            import traceback
            traceback.print_exc()
    
    #广播
    async def broadcast(self,message:str):
        for connection in self.active_connections:
            await connection.send_text(message)
manager=ConnectionManager()